#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include "tbf.h"

/*
 * 令牌桶结构体
 * token: 当前令牌数量
 * cps: 每秒生成的令牌数量 (Tokens per second)
 * burst: 令牌桶最大容量
 * mutex: 每个桶的独立互斥锁
 * cond: 条件变量
 */
typedef struct {
    int token;
    int cps;
    int burst;
    pthread_mutex_t mutex;
    pthread_cond_t cond;
} tbf_t;


static tbf_t *tbf_lib[MAXTBF]; // 令牌桶实例库，存储所有创建的令牌桶指针
static pthread_mutex_t mut_lib = PTHREAD_MUTEX_INITIALIZER;
static pthread_once_t once_control = PTHREAD_ONCE_INIT;

static int cur_tbf;     // 当前库中有效的令牌桶数量
static pthread_t tid;   // 线程id

/*
 * 线程处理函数
 * 每秒触发一次，为所有令牌桶添加令牌
 * @arg: 参数
 */
static void *__thr_handler(void *arg)
{
    int i;

    while (1) {
        pthread_mutex_lock(&mut_lib);
        
        // 遍历所有桶，增加令牌（不超过桶容量）
        for (i = 0; i < MAXTBF; i++) {
            if (NULL == tbf_lib[i])
                continue;

            pthread_mutex_lock(&tbf_lib[i]->mutex);

            tbf_lib[i]->token = (tbf_lib[i]->token + tbf_lib[i]->cps) > tbf_lib[i]->burst ? \
                                tbf_lib[i]->burst : tbf_lib[i]->token + tbf_lib[i]->cps;

            pthread_cond_broadcast(&tbf_lib[i]->cond); // 通知

            pthread_mutex_unlock(&tbf_lib[i]->mutex);
        }

        pthread_mutex_unlock(&mut_lib);
        sleep(1);
    }
}

/*
 * 模块初始化函数
 */
static void __init_moduler(void)
{
	int err;

	err = pthread_create(&tid, NULL, __thr_handler, NULL);
	if (err) {
		fprintf(stderr, "pthread_create():%s", strerror(err));
		return;
	}
}

/*
 * 初始化令牌桶
 * @cps: 令牌生成速率 (个/秒)
 * @burst: 桶容量上限
 * 返回值: 成功返回桶ID(数组下标)，失败返回-1
 */
int tbf_init(int cps, int burst)
{
    tbf_t *tbf = NULL;
    int index, err;

    // 参数校验与容量检查
    if (cps <= 0 || burst < 1 || cur_tbf == MAXTBF)
        return -1;

    // 创建令牌桶实例
    tbf = (tbf_t *)malloc(sizeof(tbf_t));
    if (NULL == tbf)
        return -1;

    // 初始化桶属性
    tbf->token = 0;     // 初始令牌为0
    tbf->cps = cps;
    tbf->burst = burst;
    
    // 初始化互斥锁
    if ((err = pthread_mutex_init(&tbf->mutex, NULL)) != 0) {
        free(tbf);
        fprintf(stderr, "pthread_mutex_init():%s\n", strerror(err));
        return -1;
    }
    // 初始化条件变量
    if ((err = pthread_cond_init(&tbf->cond, NULL)) != 0) {
        free(tbf);
        fprintf(stderr, "pthread_cond_init():%s\n", strerror(err));
        return -1;
    }

    // 寻找空槽位存储桶指针
    pthread_mutex_lock(&mut_lib);
    for (index = 0; index < MAXTBF; index++) {
        if (NULL == tbf_lib[index]) {
            tbf_lib[index] = tbf;
            cur_tbf++;  // 更新有效桶计数
            break;
        }
    }
    pthread_mutex_unlock(&mut_lib);

    // 只有首次初始化时执行
    pthread_once(&once_control, __init_moduler);

    return index;       // 返回桶ID
}

/*
 * 从令牌桶获取令牌
 * @td: 令牌桶ID
 * @ntoken: 请求的令牌数量
 * 返回值: 实际获取的令牌数（可能小于请求值），错误返回-1
 */
int tbf_fetch_token(int td, int ntoken)
{
    int ret = 0;
    tbf_t *bucket = NULL;

    // 参数校验
    if (td < 0 || td >= MAXTBF || ntoken < 1 || NULL == tbf_lib[td])
        return -1;

    bucket = tbf_lib[td];

    pthread_mutex_lock(&bucket->mutex);

    // 无令牌时阻塞等待信号添加
    while (bucket->token <= 0)
        pthread_cond_wait(&bucket->cond, &bucket->mutex);

    // 分配令牌（不超过当前可用量）
    if (bucket->token >= ntoken) {
        bucket->token -= ntoken;
        ret = ntoken;
    } else {
        ret = bucket->token;  // 返回剩余全部令牌
        bucket->token = 0;
    }

    pthread_mutex_unlock(&bucket->mutex);

    return ret;
}

/*
 * 归还令牌到桶中（不超过burst限制）
 * @td: 令牌桶ID
 * @ntoken: 归还的令牌数量
 * 返回值: 成功0，失败-1
 */
int tbf_return_token(int td, int ntoken)
{
    // 参数校验（包括归还数量不超过容量）
    if (td < 0 || td >= MAXTBF || ntoken < 1 || \
            NULL == tbf_lib[td] || ntoken > tbf_lib[td]->burst)
        return -1;

    pthread_mutex_lock(&tbf_lib[td]->mutex);
    // 增加令牌（不超过容量上限）
    tbf_lib[td]->token = (tbf_lib[td]->token + ntoken) > tbf_lib[td]->burst ? \
                         tbf_lib[td]->burst : tbf_lib[td]->token + ntoken;

    pthread_cond_broadcast(&tbf_lib[td]->cond);
    pthread_mutex_unlock(&tbf_lib[td]->mutex);
    
    return 0;
}

/*
 * 销毁指定令牌桶
 * @td: 令牌桶ID
 */
void tbf_destroy(int td)
{
    // 参数校验
    if (td < 0 || td >= MAXTBF || NULL == tbf_lib[td])
        return;

    pthread_mutex_lock(&mut_lib);

    pthread_mutex_destroy(&tbf_lib[td]->mutex);
	pthread_cond_destroy(&tbf_lib[td]->cond);

    // 释放内存并更新指针
    free(tbf_lib[td]);
    tbf_lib[td] = NULL;
    cur_tbf--;  // 更新有效桶计数
    
    pthread_mutex_unlock(&mut_lib);
}

/*
 * 销毁所有令牌桶并释放资源
 */
void tbf_destroy_all(void)
{
    int i;

    for (i = 0; i < MAXTBF; i++)
        tbf_destroy(i);

    pthread_cancel(tid); // 取消线程
    pthread_mutex_destroy(&mut_lib);
}

